import torch
import numpy as np


# 必须具有preds 和 labels
class BinaryMetric(object):
    def __init__(self, preds, labels, out_key, label_key, smooth=1):
        """
        获得二分类的混淆矩阵并计算相应的统计指标
        混淆矩阵中：TP-1 FP-2 FN-3 TN-4
        Args:
            pre: 预测值的全部输出
            label: 标签值的全部输出
        """
        super(BinaryMetric, self).__init__()
        pre, label = preds[out_key], labels[label_key]
        label = label.to(pre.device)
        pre = torch.clamp(pre, 0, 1).to(torch.int8)
        label = torch.clamp(label, 0, 1).to(torch.int8)
        self.smooth = smooth
        self.ConfusionMatrix = torch.zeros_like(pre)
        self.TP = pre * label
        self.FP = pre * (1-label)
        self.FN = (1-pre) * label
        self.TN = (1-pre) * (1-label)
        self.ConfusionMatrix = self.TP * 1 + self.FP * 2 + self.FN * 3 + self.TN * 4

    def get_confusion(self):
        return self.ConfusionMatrix

    def get_metric(self):
        tp_num, fp_num, fn_num, tn_num = torch.sum(self.TP), torch.sum(self.FP), torch.sum(self.FN), torch.sum(self.TN)
        if tp_num + fn_num == 0:
            Precision = 0
        else:
            Precision = tp_num / (tp_num + fp_num)

        if tp_num + fn_num == 0:
            Recall = 0
        else:
            Recall = tp_num / (tp_num + fn_num)

        F1 = (2 * tp_num) / (2 * tp_num + fn_num + fp_num)

        return float(Precision), float(Recall), float(F1)

    def get_num(self):
        tp_num, fp_num, fn_num, tn_num = torch.sum(self.TP), torch.sum(self.FP), torch.sum(self.FN), torch.sum(self.TN)
        return tp_num, fp_num, fn_num, tn_num

    def __call__(self, *args, **kwargs):
        return self.get_metric()